# uncompyle6 version 3.2.3
# Python bytecode 3.6 (3379)
# Decompiled from: Python 3.6.8 |Anaconda custom (64-bit)| (default, Feb 21 2019, 18:30:04) [MSC v.1916 64 bit (AMD64)]
# Embedded file name: helpers\rate_limiter\leacky_bucket_collection.py
__author__ = "Attila Gerendi (Sanyi)"
import threading, time, bisect

IX_WILL_EMPTY = 0
IX_TOKENS = 1
IX_LAST_LEAK = 2


class LeakyBucketCollection(object):
    def __init__(self, leak_rate, max_tokens, max_buckets=100, time_function=time.time):
        self._rate = leak_rate
        self._size = max_tokens
        self._max_buckets = max_buckets
        self._time = time_function
        self.buckets = {}
        self._buckets_sorting = []
        self._lock = threading.Lock()

    def get_tokens(self, key):
        """
        returns the number of tokens in bucket
        :param key:
        :return:
        """
        ct = self._time()
        with self._lock:
            bucket = self.buckets.get(key)
        if bucket is None:
            return 0
        leaked = (ct - bucket[IX_LAST_LEAK]) * self._rate
        tokens = bucket[IX_TOKENS] - leaked
        if tokens > 0:
            return tokens
        else:
            return 0

    def is_full(self, key):
        """
        returns true if tokens in the bucket exceed threshold
        :param key:
        :return:
        """
        ct = self._time()
        with self._lock:
            bucket = self.buckets.get(key)
        if bucket is None:
            return False
        else:
            leaked = (ct - bucket[IX_LAST_LEAK]) * self._rate
            tokens = bucket[IX_TOKENS] - leaked
            return tokens > self._size

    def wait_time(self, key, till_empty=False):
        """
        returns wait time till threshold or bucket is empty
        :param key:
        :param till_empty:
        :return:
        """
        ct = self._time()
        with self._lock:
            bucket = self.buckets.get(key)
        if bucket is None:
            return 0
        leaked = (ct - bucket[IX_LAST_LEAK]) * self._rate
        tokens = bucket[IX_TOKENS] - leaked
        if tokens < 0:
            return 0
        elif till_empty:
            return tokens / self._rate
        elif tokens <= self._size:
            return 0
        else:
            return (tokens - self._size) / self._rate

    def add(self, key, tokens=1):
        """
        Returns time necessary for the bucket to leak under threshold (max_tokens)
        :param key:
        :param tokens:
        :return:
        """
        with self._lock:
            ct = self._time()
            bucket = self.buckets.get(key)
            if bucket is None:
                will_empty = (ct + tokens) / self._rate
            else:
                i = bisect.bisect_left(
                    self._buckets_sorting, (bucket[IX_WILL_EMPTY], key)
                )
                if not i != len(self._buckets_sorting):
                    raise AssertionError
                if not self._buckets_sorting[i] == (bucket[IX_WILL_EMPTY], key):
                    raise AssertionError
                del self._buckets_sorting[i]
                leaked = (ct - bucket[IX_LAST_LEAK]) * self._rate
                tokens_before = bucket[IX_TOKENS] - leaked
                if tokens_before < 0:
                    tokens_before = 0
                tokens += tokens_before
                will_empty = (ct + tokens) / self._rate
            self.buckets[key] = (will_empty, tokens, ct)
            bisect.insort(self._buckets_sorting, (will_empty, key))
            i = 0
            l = len(self.buckets)
            while i < l:
                if self._buckets_sorting[i][0] < ct:
                    del self.buckets[self._buckets_sorting[i][1]]
                    i += 1

            if i > 0:
                self._buckets_sorting = self._buckets_sorting[i:]
            else:
                if len(self.buckets) > self._max_buckets:
                    del self.buckets[self._buckets_sorting[0][1]]
                    self._buckets_sorting = self._buckets_sorting[1:]
            print(tokens, self._size, tokens / self._rate)
            if tokens <= self._size:
                return 0
            return (tokens - self._size) / self._rate


if __name__ == "__main__":

    class FakeTime:
        def __init__(self, v=0):
            self.time = v

        def get_time(self):
            return self.time

    t = FakeTime()
    b = LeakyBucketCollection(0.1, 2, time_function=t.get_time)
    print(b.add("x", 1), b.get_tokens("x"))
    print(b.wait_time("x"), b.get_tokens("x"))
    print(b.add("x", 1), b.get_tokens("x"))
    print(b.wait_time("x"), b.get_tokens("x"))
    print(b.add("x", 1), b.get_tokens("x"))
    print(b.wait_time("x"), b.get_tokens("x"))
    t.time = 1
    print(b.wait_time("x"), b.get_tokens("x"))
    t.time = 2
    print(b.wait_time("x"), b.get_tokens("x"))
    t.time = 3
    print(b.add("x", 1), b.get_tokens("x"))
    print(b.wait_time("x"), b.get_tokens("x"))
    t.time = 10
    print(b.wait_time("x"), b.get_tokens("x"))
    t.time = 20
    print(b.wait_time("x"), b.get_tokens("x"))
